{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "onzVg96bnvsD"
      },
      "outputs": [],
      "source": [
        "#@title 👇 安装 Waifuc\n",
        "print(\"安装 Waifuc 中\")\n",
        "%pip install huggingface_hub git+https://github.com/deepghs/waifuc.git@main#egg=waifuc\n",
        "print(\"安装 Waifuc 完成\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "z53kDygQoUmZ"
      },
      "outputs": [],
      "source": [
        "#@title 👇 爬取图片\n",
        "import os\n",
        "from waifuc.action import (\n",
        "    HeadCountAction,\n",
        "    NoMonochromeAction,\n",
        "    FilterSimilarAction,\n",
        "    TaggingAction,\n",
        "    PersonSplitAction,\n",
        "    FaceCountAction,\n",
        "    FirstNSelectAction,\n",
        "    CCIPAction,\n",
        "    ModeConvertAction,\n",
        "    ClassFilterAction,\n",
        "    RandomFilenameAction,\n",
        "    AlignMinSizeAction\n",
        ")\n",
        "from waifuc.export import SaveExporter\n",
        "from waifuc.source import DanbooruSource, YandeSource\n",
        "\n",
        "tag = \"nikorashi-ka\" #@param {type:\"string\"}\n",
        "count = \"5\" #@param {type:\"string\"}\n",
        "path = \"/content/waifuc\" #@param {type:\"string\"}\n",
        "resource = \"danbooru\" #@param [\"danbooru\", \"yande.re\"]\n",
        "\n",
        "print(f\"开始爬取图片, 爬取标签: {tag}, 保存路径: {path}\")\n",
        "\n",
        "if resource == \"danbooru\":\n",
        "    source = DanbooruSource([tag])\n",
        "elif resource == \"yande.re\":\n",
        "    source = YandeSource([tag])\n",
        "\n",
        "source.attach(\n",
        "    HeadCountAction(1), # only 1 head,\n",
        "    ModeConvertAction('RGB', 'white'), # 以RGB色彩模式加载图像并将透明背景替换为白色背景\n",
        "    ClassFilterAction(['illustration', 'bangumi']),  # 丢弃漫画或3D图像\n",
        "    FilterSimilarAction('all'),  # 丢弃相似或重复的图像\n",
        ")[:int(count)].export(\n",
        "    SaveExporter(os.path.join(path, f\"waifuc_{tag}\"))\n",
        ")\n",
        "\n",
        "print(f\"爬取图片完成, 保存路径: {path}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "C2bSvuZqPkSB"
      },
      "outputs": [],
      "source": [
        "#@title 👇 上传爬取的数据到 HuggingFace\n",
        "import os\n",
        "from huggingface_hub import HfApi\n",
        "\n",
        "\n",
        "api = HfApi()\n",
        "repo_id = \"\" #@param {type:\"string\"}\n",
        "repo_type = \"dataset\" #@param [\"model\", \"dataset\", \"space\"]\n",
        "upload_path = \"/content/waifuc\" #@param {type:\"string\"}\n",
        "hf_token = \"\" #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "print(f\"上传 {upload_path} 到 {repo_id} (类型: {repo_type}) 中\")\n",
        "os.environ[\"HF_TOKEN\"] = hf_token\n",
        "try:\n",
        "    api.upload_folder(\n",
        "        repo_id=repo_id,\n",
        "        repo_type=repo_type,\n",
        "        folder_path=upload_path\n",
        "    )\n",
        "except:\n",
        "    api.upload_large_folder(\n",
        "        repo_id=repo_id,\n",
        "        repo_type=repo_type,\n",
        "        folder_path=upload_path\n",
        "    )\n",
        "print(f\"上传 {upload_path} 到 {repo_id} (类型: {repo_type}) 完成\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "欢迎使用 Colaboratory",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.11.8"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
